# ICEMGenerator - 高效的ICEM CFD网格划分Replay Script生成工具
# 模块： 
# (1) Curve类，是所有创建线（含直线、圆环、弧线）的类的超类
# (2) 在其基础上衍生出的直线类Line、圆环类Ring和弧线类Arc
# 作者： weixc
# 邮箱： 12028045 [at] zju.edu.cn
# 项目主页： https://gitee.com/weixinchong/icemgenerator

import abc, numpy as np
from collections.abc import Iterable
from .point import Point, create_points
from .param_check import ListSizeError, check_axis

def sind(deg: float) -> float:
    '''角度正弦函数
    :param deg: 角度
    :return: 正弦值'''
    return np.sin(deg / 180 * np.pi)

def cosd(deg: float) -> float:
    '''角度余弦函数
    :param deg: 角度
    :return: 余弦值'''
    return np.cos(deg / 180 * np.pi)

class Curve():
    '''所有线的超类'''

    curr_cId = 0  #类静态变量，用于确定最新创建线的编号
    members = []  #现有线的成员列表

    def init_cId(cId: int):
        '''初始化最新创建线的编号
        :param cId: 创建的第一条线的编号'''
        Curve.curr_cId = cId

    def members_replay(assisted: bool = False) -> str:
        '''生成所有线的Replay Script
        :param assisted: 要生成普通曲线还是辅助线的Replay Script
        :return: 生成的Replay Script'''
        text = ''
        for c in Curve.members:
            if c.assisted == assisted:
                text += c.replay() + '\n'
        return text

    def __str__(self):
        return 'crv.%03d' % self.cId

    def __repr__(self):
        return self.__str__()

    def setName(self, name: str):
        '''更改线的名称
        :param name: 要设定的线的名称'''
        self.name = name

    @abc.abstractmethod
    def replay(self):
        '''生成线的Replay Script'''
        pass

    def info(self, single: bool = False) -> str:
        '''获得线的信息
        :param single: 是否输出单一信息，为True时仅输出线的信息，为False时还会输出每个点的信息
        :return: 线的信息'''
        # 先输出线的信息
        text = 'Curve %03d' % self.cId
        if self.name != None:
            text += ' [%s]' % self.name
        if self.type == 'line':
            text += ': %s - %s' % (self.points[0], self.points[1])
        elif self.type == 'ring':
            text += ': %s - (%s, %s)' % (self.points[0], self.points[1], self.points[2])
        elif self.type == 'arc':
            text += ': %s -- %s' % (self.points[0], self.points[-1])
        if not single:
            # 再输出点的信息
            text += '\n[Corresponding Points]\n%s\n%s' % (self.points[0].info(), self.points[1].info())
        return text

    def getPoint(self, name: str) -> Point | None:
        '''获得线上指定名称的点
        :param name: 点的名称
        :return: 找到指定名称的点时，返回该点；否则返回None'''
        for p in self.points:
            if p.name == name:
                return p


class Line(Curve):
    '''创建一条直线'''

    def find_member(p1: Point, p2: Point) -> Curve | None:
        '''查找指定的直线是否已经定义过
        :params p1, p2: 连成要查找的直线的点，p1、p2的顺序不影响查找结果
        :return (Line | None): 找到已定义的直线时，返回该直线；否则为None'''
        for c in Curve.members:
            if c.type == 'line':
                if (c.points[0] is p1 and c.points[1] is p2) or (c.points[1] is p1 and c.points[0] is p2):
                    return c

    def create(p1: Point, p2: Point, name: str = None, assisted: bool = False) -> Curve:
        '''创建一条直线（已存在则引用，不存在则创建）
        :params p1, p2: 连成直线的两个点
        :param name: 直线的名称，非必须
        :param assisted: 表明直线为划分Block使用的辅助线，辅助线将于面创建完成后创建
        :return (Line): 创建的直线'''
        crv = Line.find_member(p1, p2)
        if crv == None:
            crv = Line(p1, p2, name, assisted)
        elif name != None:
            crv.setName(name)
        return crv

    def __init__(self, p1, p2, name, assisted):
        '''创建一条直线'''
        self.cId = Curve.curr_cId
        self.points = (p1, p2)
        self.direction = p1.X - p2.X
        self.name = name
        self.assisted = assisted
        self.type = 'line'
        Curve.members.append(self)  #将当前线增加到成员列表
        Curve.curr_cId += 1  #最新线编号增加1

    def replay(self) -> str:
        '''生成直线的Replay Script
        :return: 生成的Replay Script'''
        return 'ic_curve point GEOM %s {%s %s}' % (self, self.points[0], self.points[1])

class Ring(Curve):
    '''创建一条圆环'''

    def find_member(r: float, x: float, y: float, z: float, axis: str) -> Curve | None:
        '''查找指定的圆环是否已经定义过
        :param r: 圆环的半径
        :param x, y, z: 圆环的圆心
        :param axis: 弧线的轴向，可选值为x、y、z
        :return (Ring): 找到已定义的圆环时，返回该圆环；否则为None'''
        check_axis(axis)
        for c in Curve.members:
            if c.type == 'ring':
                if c.r == r and c.x == x and c.y == y and c.z == z and c.axis == axis:
                    return c

    def create(r: float, x: float = 0, y: float = 0, z: float = 0, axis: str = 'z', name: str = None, assisted: bool = False) -> Curve:
        '''创建一个圆环（已存在则引用，不存在则创建）
        :param r: 圆环的半径
        :params x, y, z: 圆环的圆心
        :param axis: 弧线的轴向，可选值为x、y、z
        :param name: 弧线的名称，非必须
        :param assisted: 表明曲线为划分Block使用的辅助线，辅助线将于面创建完成后创建
        :return (Ring): 创建的圆环'''
        check_axis(axis)
        crv = Ring.find_member(r, x, y, z, axis)
        if crv == None:
            crv = Ring(r, x, y, z, axis, name, assisted)
        elif name != None:
            crv.setName(name)
        return crv

    def __init__(self, r, x, y, z, axis, name, assisted):
        '''创建一个圆环'''
        self.cId = Curve.curr_cId
        self.r = r
        self.x = x
        self.y = y
        self.z = z
        self.axis = axis
        # 创建环的各个点
        pnt1 = Point.create(x, y, z)
        if axis == 'x':
            pnt2 = Point.create(x, y + r, z)
            pnt3 = Point.create(x, y, z + r)
            self.direction = np.array([1,0,0])
        elif axis == 'y':
            pnt2 = Point.create(x + r, y, z)
            pnt3 = Point.create(x, y, z + r)
            self.direction = np.array([0,1,0])
        elif axis == 'z':
            pnt2 = Point.create(x + r, y, z)
            pnt3 = Point.create(x, y + r, z)
            self.direction = np.array([0,0,1])
        self.points = [pnt1, pnt2, pnt3]
        self.name = name
        self.assisted = assisted
        self.type = 'ring'
        Curve.members.append(self)  #将当前线增加到成员列表
        Curve.curr_cId += 1  #最新线编号增加1

    def replay(self) -> str:
        '''生成圆环的Replay Script
        :return: 生成的Replay Script'''
        return 'ic_curve arc_ctr_rad GEOM %s {%s %s %s %g 0 360}' % (self, self.points[0], self.points[1], self.points[2], self.r)


class Arc(Curve):
    '''创建一条弧线'''

    def find_member(r: float, x: float, y: float, z: float, axis: str, angle_start: float, angle_end: float, axis_offset: float) -> Curve | None:
        '''查找指定的弧线是否已经定义过
        :param r: 弧线的半径
        :param x, y, z: 弧线的圆心
        :param axis: 弧线的轴向，可选值为x、y、z
        :param angle_start, angle_end: 表示弧线的角度从angle_start_t到angle_end_t
        :param axis_offset: 弧线沿轴向的上升高度，当axis_offset取负时，表示弧线沿负半轴下降
        :return: 找到已定义的弧线时，返回该弧线；否则为None'''
        check_axis(axis)
        for c in Curve.members:
            if c.type == 'arc':
                if c.r == r and c.x == x and c.y == y and c.z == z and c.axis == axis and (c.angle_start == angle_start and c.angle_end == angle_end or c.angle_start == angle_end and c.angle_end == angle_start) and c.axis_offset == axis_offset:
                    return c

    def create(r: float, x: float = 0, y: float = 0, z: float = 0, axis: str = 'z', \
               angle_start: float = 0, angle_end: float = 90, axis_offset: float = 0, name: str = None, assisted: bool = False) -> Curve:
        '''创建一段弧线（已存在则引用，不存在则创建）
        :param r: 弧线的半径
        :param x, y, z: 弧线的圆心
        :param axis: 弧线的轴向，可选值为x、y、z
        :param angle_start, angle_end: 表示弧线的角度从angle_start_t到angle_end_t
        :param axis_offset: 弧线沿轴向的上升高度，当axis_offset取负时，表示弧线沿负半轴下降
        :param name: 弧线的名称，非必须
        :param assisted: 表明曲线为划分Block使用的辅助线，辅助线将于面创建完成后创建
        :return: 创建的弧线'''
        check_axis(axis)
        crv = Arc.find_member(r, x, y, z, axis, angle_start, angle_end, axis_offset)
        if crv == None:
            crv = Arc(r, x, y, z, axis, angle_start, angle_end, axis_offset, name, assisted)
        elif name != None:
            crv.setName(name)
        return crv

    def __init__(self, r, x, y, z, axis, angle_start, angle_end, axis_offset, name, assisted):
        '''创建一段弧线'''
        self.cId = Curve.curr_cId
        self.r = r
        self.x = x
        self.y = y
        self.z = z
        self.axis = axis
        self.axis_offset = axis_offset
        self.angle_start = angle_start
        self.angle_end = angle_end
        if angle_start > angle_end:  #交换角度顺序
            angle_end, angle_start = angle_start, angle_end
        # 判断角度差，当角度差>180时，需要创建多条弧线并最终合并成一条大弧
        delta = angle_end - angle_start  #总角度差
        if delta <= 180:
            self.cCount = 1
            pCount = 3  #一条弧线由三个点连成
            angles_list = [angle_start, angle_start + delta / 2, angle_end]
        else:
            self.cCount = int(delta / 180)  #弧线段数向下取整
            pCount = 1 + 2 * self.cCount  #每增加一条弧线，增加2个点
            delta_each = delta / (pCount - 1)  #每两点之间的角度差
            angles_list = [None] * pCount
            angles_list[0] = angle_start
            for i in range(1, pCount):
                angles_list[i] = angles_list[i - 1] + delta_each
        axis_offset_each = axis_offset / (pCount - 1)
        # 创建环的各个点
        pnt = []
        a_offset = 0
        if axis == 'x':
            for a in angles_list:
                pnt.append(Point.create(x + a_offset, y + r * cosd(a), z + r * sind(a)))
                a_offset += axis_offset_each
            self.direction = np.array([1,0,0])
        elif axis == 'y':
            for a in angles_list:
                pnt.append(Point.create(x + r * cosd(a), y + a_offset, z + r * sind(a)))
                a_offset += axis_offset_each
            self.direction = np.array([0,1,0])
        elif axis == 'z':
            for a in angles_list:
                pnt.append(Point.create(x + r * cosd(a), y + r * sind(a), z + a_offset))
                a_offset += axis_offset_each
            self.direction = np.array([0,0,1])
        self.points = pnt
        self.name = name
        self.assisted = assisted
        self.type = 'arc'
        Curve.members.append(self)  #将当前线增加到成员列表
        if self.cCount == 1:
            Curve.curr_cId += 1  #不需要将多条线合并时，最新线编号增加1
        else:
            Curve.curr_cId += self.cCount + 1  #不需要将多条线合并时，最新线编号在增加1的基础上还需要增加弧线段数

    def __str__(self):
        if self.cCount == 1:
            return 'crv.%03d' % self.cId
        else:
            return 'crv.%03d' % (self.cId + self.cCount)

    def replay(self) -> str:
        '''生成弧线的Replay Script
        :return: 生成的Replay Script'''
        if self.cCount == 1:  #单条弧线可直接创建
            return 'ic_curve arc GEOM %s {%s %s %s}' % (self, self.points[0], self.points[1], self.points[2])
        else:  #多条弧线需要先分别创建，然后合并成一条大弧
            text = ''
            crv_list = [None] * self.cCount
            for i in range(self.cCount):
                crv_list[i] = 'crv.%03d' % (self.cId + i)
                text += 'ic_curve arc GEOM %s {%s %s %s}\n' % (crv_list[i], self.points[i*2], self.points[i*2+1], self.points[i*2+2])
            text += 'ic_curve concat GEOM %s {%s}' % (self, ' '.join(crv_list))
            return text


class Arcp(Curve):
    '''基于三点创建一条弧线'''

    def find_member(pnt1: Point, pnt2: Point, pnt3: Point) -> Curve | None:
        '''查找指定的弧线是否已经定义过
        :param pnt1, pnt2, pnt3: 弧线上的相邻三个点
        :return: 找到已定义的弧线时，返回该弧线；否则为None'''
        for c in Curve.members:
            if c.type == 'arcp':
                if c.points[0] == pnt1 and c.points[1] == pnt2 and c.points[2] == pnt3:
                    return c
                
    def create(pnt1: Point, pnt2: Point, pnt3: Point, name: str = None, assisted: bool = False) -> Curve:
        '''创建一段三点连成的弧线（已存在则引用，不存在则创建）
        :param pnt1, pnt2, pnt3: 弧线上的相邻三个点
        :param name: 弧线的名称，非必须
        :param assisted: 表明曲线为划分Block使用的辅助线，辅助线将于面创建完成后创建
        :return: 创建的弧线'''
        crv = Arcp.find_member(pnt1, pnt2, pnt3)
        if crv == None:
            crv = Arcp(pnt1, pnt2, pnt3, name, assisted)
        elif name != None:
            crv.setName(name)
        return crv
    
    def __init__(self, pnt1: Point, pnt2: Point, pnt3: Point, name, assisted):
        '''创建一段三点连成的弧线'''
        self.cId = Curve.curr_cId
        self.points = [pnt1, pnt2, pnt3]
        self.direction = pnt3.X - pnt1.X  #沿弧线首尾连一条线，视为法向量
        self.name = name
        self.assisted = assisted
        self.type = 'arcp'
        Curve.members.append(self)  #将当前线增加到成员列表
        Curve.curr_cId += 1  #最新线编号增加1

    def replay(self) -> str:
        '''生成弧线的Replay Script
        :return: 生成的Replay Script'''
        return 'ic_curve arc GEOM %s {%s %s %s}' % (self, self.points[0], self.points[1], self.points[2])


def lines_from_points(points: Iterable | tuple) -> tuple:
    '''由一系列点创建连续的直线
    :param points: 一系列点，连线顺序为points[0]-points[1]-points[2]-..-points[n]
    :returns: 由n点连成的n-1条直线，其中第i条直线由points[i]与points[i+1]连成'''
    curves = []
    for i in range(len(points) - 1):
        curves.append(Line.create(points[i], points[i + 1]))
    return tuple(curves)

def create_line(x: float | Iterable, y: float | Iterable, z: float | Iterable = None, name: str = None) -> Line:
    '''根据端点坐标值创建直线
    :params x, y, z: 对于x、y、z，其中不少于1个变量为列表，对于二维模型，参数`z`可以为空，例如:
        (1)x=(1,3), y=5: 直线由(1,5,0)和(3,5,0)连成
        (2)x=(1,3), y=(5,7), z=(3,9): 直线由(1,5,3)连成(3,7,9)
    :param name: 直线的名称
    :return: 连成的直线'''
    if not isinstance(x, Iterable):
        x = (x, x)
    if not isinstance(y, Iterable):
        y = (y, y)
    if not isinstance(z, Iterable):
        z = (z, z)
    if len(x) > 2 or len(y) > 2 or len(z) > 2:
        raise ListSizeError('创建直线时，参数x,y,z最多只能是两个参数构成的列表')
    if x[0] == x[1] and y[0] == y[1] and z[0] == z[1]:
        raise ListSizeError('创建直线时，参数x,y,z至少有一个是列表')
    pnt1 = Point.create(x[0], y[0], z[0])
    pnt2 = Point.create(x[1], y[1], z[1])
    return Line.create(pnt1, pnt2, name)

def create_lines(x: Iterable, y: Iterable, z: Iterable = None) -> list:
    '''根据端点坐标值创建连续的直线
    :params x, y, z: 各点的端点坐标值构成的列表，对于二维模型，参数`z`可以为空，例如:
        (1)x=(0,2,2,0),y=(0,0,3,0): 创建了(0,0,0)、(2,0,0)、(2,3,0)、(0,0,0)相连形成的直线
        (2)x=(0,2,2,0),y=(0,0,3,0),z=(0,1,4,0): 创建了(0,0,0)、(2,0,1)、(2,3,4)、(0,0,0)相连形成的直线
    :param name: 直线的名称
    :return: 连成的直线'''
    if z is None:
        z = [0] * len(x)
    if len(x) != len(y) or len(x) != len(z) or len(y) != len(z):
        raise ListSizeError('创建连续直线时，参数x,y,z的元素个数应保持一致')
    N_lines = len(x) - 1
    lines = [None] * N_lines
    for i in range(N_lines):
        pnt1 = Point.create(x[i], y[i], z[i])
        pnt2 = Point.create(x[i+1], y[i+1], z[i+1])
        lines[i] = Line.create(pnt1, pnt2)
    return lines

def create_crv_by_offset(crv_old: Curve, offset: float, axis: str = 'z', name: str = None) -> Curve:
    '''通过对已有曲线的偏移，创建新曲线
    :param crv_old: 待偏移的曲线
    :param offset: 沿axis方向的偏移量
    :param axis: 偏移方向
    :param name: 偏移后曲线的名称
    :return: 偏移后的新曲线'''
    if name is None:
        name = f'_{axis}offset_{offset}'
    if type(crv_old) == Line:  #直线偏移
        p1, p2 = crv_old.points
        p1_new = Point.create(
            p1.x + (offset if axis == 'x' else 0), 
            p1.y + (offset if axis == 'y' else 0), 
            p1.z + (offset if axis == 'z' else 0), 
            p1.name)
        p2_new = Point.create(
            p2.x + (offset if axis == 'x' else 0), 
            p2.y + (offset if axis == 'y' else 0), 
            p2.z + (offset if axis == 'z' else 0), 
            p2.name)
        crv_new = Line.create(p1_new, p2_new, name)
    elif type(crv_old) == Ring:  #圆环偏移
        crv_new = Ring.create(
            crv_old.r, 
            crv_old.x + (offset if axis == 'x' else 0), 
            crv_old.y + (offset if axis == 'y' else 0), 
            crv_old.z + (offset if axis == 'z' else 0), 
            crv_old.axis, name)
    elif type(crv_old) == Arc:  #圆弧偏移
        crv_new = Arc.create(
            crv_old.r, 
            crv_old.x + (offset if axis == 'x' else 0), 
            crv_old.y + (offset if axis == 'y' else 0), 
            crv_old.z + (offset if axis == 'z' else 0), 
            crv_old.axis, crv_old.angle_start, crv_old.angle_end, 
            crv_old.axis_offset, name)
    return crv_new

def create_crvs_by_offset(crvs_old: Iterable, offset: float, axis: str = 'z', suffix: str = None) -> tuple:
    '''通过对一系列已有曲线的偏移，创建新曲线
    :param crvs_old: 待偏移的曲线列表
    :param offset: 沿axis方向的偏移量
    :param axis: 偏移方向
    :param suffix: 偏移后曲线名称的后缀
    :return: 偏移后的新曲线列表'''
    crvs_new = []
    for crv in crvs_old:  #逐条曲线进行偏移
        if suffix is None:
            name = crv.name + f'_{axis}offset_{offset}'
        else:
            name = crv.name + suffix
        crv_new = create_crv_by_offset(crv, offset, axis, name)
        crvs_new.append(crv_new)
    return tuple(crv_new)